import os
import comtypes.client
from comtypes.gen import STKObjects, STKUtil
from abc import ABC, abstractmethod

class Object(ABC):
    """对象抽象基类"""
    @abstractmethod
    def create(self, scenario, name, config):
        pass

class AttachedObject(ABC):
    """附加对象抽象基类"""
    @abstractmethod
    def attach(self, parent, sensor_config, root):
        pass
            
class Satellite(Object):
    """卫星类"""
    def create(self, scenario, name, obj_config):
        """创建卫星对象
        Args:
            scenario: STK场景对象
            name (str): 卫星名称
            obj_config (dict): 卫星配置字典
        """
        try:
            if not obj_config or obj_config['type'] != 'satellite':
                raise ValueError(f"无效的卫星配置: {name}")

            # 获取轨道参数
            orbit_config = obj_config['basic']['orbit']

            # 创建卫星对象
            satellite = scenario.Children.New(STKObjects.eSatellite, name)
            satellite2 = satellite.QueryInterface(STKObjects.IAgSatellite)

            # 设置传播器
            propagator_type = getattr(STKObjects, orbit_config['propagator'])

            if propagator_type == STKObjects.ePropagatorTwoBody:
                satellite2.SetPropagatorType(propagator_type)
                propagator = satellite2.Propagator.QueryInterface(STKObjects.IAgVePropagatorTwoBody)
                # 设置轨道参数
                keplerian = propagator.InitialState.Representation.ConvertTo(
                    STKUtil.eOrbitStateClassical).QueryInterface(STKObjects.IAgOrbitStateClassical)
                
                keplerian.SizeShapeType = STKObjects.eSizeShapeSemimajorAxis
                keplerian.SizeShape.QueryInterface(
                    STKObjects.IAgClassicalSizeShapeSemimajorAxis).SemiMajorAxis = orbit_config['semi_major_axis']
                keplerian.SizeShape.QueryInterface(
                    STKObjects.IAgClassicalSizeShapeSemimajorAxis).Eccentricity = orbit_config['eccentricity']
                # 倾角是指轨道面与参考平面（通常是赤道平面）之间的角度
                keplerian.Orientation.Inclination = orbit_config['inclination']
                # 近地点辐角
                keplerian.Orientation.ArgOfPerigee = orbit_config['arg_perigee']
                # 设置Kepler轨道参数中的升交点类型为升交点赤经（RAAN）
                keplerian.Orientation.AscNodeType = STKObjects.eAscNodeRAAN
                keplerian.Orientation.AscNode.QueryInterface(
                    STKObjects.IAgOrientationAscNodeRAAN).Value = orbit_config['raan']
                
                # 设置keplerian对象的位置类型为真近点角（True Anomaly表示从近地点（轨道上离中心天体最近点）到当前点的角度）
                keplerian.LocationType = STKObjects.eLocationTrueAnomaly
                # 通过QueryInterface方法获取keplerian对象的真实近点角位置接口
                # STKObjects.IAgClassicalLocationTrueAnomaly是一个接口，用于处理经典轨道位置中的真近点角
                keplerian.Location.QueryInterface(
                    STKObjects.IAgClassicalLocationTrueAnomaly).Value = orbit_config['true_anomaly']

                # 应用参数并传播
                propagator.InitialState.Representation.Assign(keplerian)
                propagator.Propagate()
            
            elif propagator_type == STKObjects.ePropagatorSGP4:
                satellite2.SetPropagatorType(propagator_type)
                propagator = satellite2.Propagator.QueryInterface(propagator_type)
                # TODO:

            return satellite2

        except Exception as e:
            print(f"创建卫星失败: {str(e)}")
            raise

class Vessel(Object):
    """船舰类"""
    def create(self, scenario, name, obj_config):
        """创建船舰对象
        Args:
            scenario: STK场景对象
            name (str): 船舰名称
            obj_config (dict): 船舰配置字典
        """
        try:
            if not obj_config or obj_config['type'] != 'vessel':
                raise ValueError(f"无效的船舰配置: {name}")

            # 获取路径参数
            route_config = obj_config['basic']['route']

            # 创建船舰对象
            vessel = scenario.Children.New(STKObjects.eShip, name)
            vessel2 = vessel.QueryInterface(STKObjects.IAgShip)

            # 设置路径
            vessel2.SetRouteType(STKObjects.ePropagatorGreatArc)
            route = vessel2.Route.QueryInterface(STKObjects.IAgVePropagatorGreatArc)
            route.Method = 0

            # 添加航点
            for wp in route_config['waypoints']:
                waypoint = route.Waypoints.Add()
                waypoint.Latitude = wp['latitude']
                waypoint.Longitude = wp['longitude']
                waypoint.Altitude = wp['altitude']
                waypoint.Speed = wp['speed']

            route.Propagate()
            return vessel2

        except Exception as e:
            print(f"创建船舰失败: {str(e)}")
            raise

class SensorSAR(AttachedObject):
    """SAR传感器类"""
    def attach(self, parent_name, sensor_config, root):
        try:
            # 获取传感器参数
            sensor_params = sensor_config['params']

            # 创建传感器
            root.ExecuteCommand(
                f'New / */Satellite/{parent_name}/Sensor {sensor_config['name']}')
            
            # 配置传感器
            cmd = f'Define */Satellite/{parent_name}/Sensor/{sensor_config['name']} SAR ' \
                f'{sensor_params["elevation_angles_min"]} ' \
                f'{sensor_params["elevation_angles_max"]} ' \
                f'{sensor_params["exclusion_angle_forward"]} ' \
                f'{sensor_params["exclusion_angle_aft"]} ' \
                'TrackParentAltitude On'
            root.ExecuteCommand(cmd)

            return True

        except Exception as e:
            print(f"添加传感器失败: {str(e)}")
            raise

class STKUtils:
    def __init__(self, scenario_name:str='MyScenario'):
        """初始化STK场景
        Args:
            scenario_name (str): 场景名称
        """
        try:
            self.stk = comtypes.client.GetActiveObject('STK11.Application')
        except:
            self.stk = comtypes.client.CreateObject('STK11.Application')
            self.stk.Visible = True
            
        self.root = self.stk.Personality2
        if not self.root.CurrentScenario:
            self.root.NewScenario(scenario_name)
        self.scenario = self.root.CurrentScenario

    def create_object(self, obj_config):
        """根据配置自动创建对象
        Args:
            obj_config (dict): 对象配置
        Returns:
            创建的对象
        """
        obj_name = obj_config['name']
        obj_type = obj_config['type']
        
        if obj_type == 'satellite':
            satellite = Satellite()
            return satellite.create(self.scenario, obj_name, obj_config)
        elif obj_type == 'vessel':
            vessel = Vessel()
            return vessel.create(self.scenario, obj_name, obj_config)
        else:
            raise ValueError(f"未知的对象类型: {obj_type}")

    def add_sensor(self, parent, sensor_config):
        """添加传感器"""
        sensor = SensorSAR()
        return sensor.attach(parent, sensor_config, root = self.root)

    def save_scenario(self, save_path):
        """保存场景"""
        try:                
            folder_path = os.path.dirname(save_path)
            if not os.path.exists(folder_path):
                os.makedirs(folder_path, exist_ok=True)

            self.root.SaveAs(save_path)
            print(f"场景已保存到: {folder_path}")
            return True

        except Exception as e:
            print(f"保存场景失败: {str(e)}")
            return False

    def close(self):
        """关闭STK连接"""
        try:
            if self.scenario and not self.scenario.Saved:
                print("警告：场景未保存")
            if self.stk and self.stk.Visible:
                self.stk.Quit()
        except Exception as e:
            print(f"关闭STK失败: {str(e)}")

    def __del__(self):
        """析构函数"""
        self.close()
